/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <set>
#include <map>
#include <stdlib.h>
#include "parcsr_matmul.h"

namespace YHAMG
{

static void CSR_MulAddMul(const CSR_Matrix& A, const CSR_Matrix& B, const CSR_Matrix& C, const CSR_Matrix& D, CSR_Matrix& W)
{
	int n = A.size[0];
	int m = B.size[1];
	int* Ap = A.rowptr;
	int* Ai = A.colind;
	double* Av = A.values;
	int* Bp = B.rowptr;
	int* Bi = B.colind;
	double* Bv = B.values;
	int* Cp = C.rowptr;
	int* Ci = C.colind;
	double* Cv = C.values;
	int* Dp = D.rowptr;
	int* Di = D.colind;
	double* Dv = D.values;

	int* Wp = new int[n + 1];
	Wp[0] = 0;
#ifdef USE_OPENMP
#pragma omp parallel
#endif
	{
		int* w = new int[m];
		for (int i = 0; i < m; ++i)
			w[i] = -1;
#ifdef USE_OPENMP
#pragma omp for schedule(guided)
#endif
		for (int i = 0; i < n; ++i)
		{
			int cnt = 0;

			for (int k = Ap[i]; k < Ap[i + 1]; ++k)
			{
				int kcol = Ai[k];
				for (int j = Bp[kcol]; j < Bp[kcol + 1]; ++j)
				{
					int jcol = Bi[j];
					if (w[jcol] != i)
					{
						w[jcol] = i;
						++cnt;
					}
				}
			}

			for (int k = Cp[i]; k < Cp[i + 1]; ++k)
			{
				int kcol = Ci[k];
				for (int j = Dp[kcol]; j < Dp[kcol + 1]; ++j)
				{
					int jcol = Di[j];
					if (w[jcol] != i)
					{
						w[jcol] = i;
						++cnt;
					}
				}
			}

			Wp[i + 1] = cnt;
		}

		delete[] w;
	}

	for (int i = 0; i < n; ++i)
		Wp[i + 1] += Wp[i];

	int* Wi = new int[Wp[n]];
	double* Wv = new double[Wp[n]];

#ifdef USE_OPENMP
#pragma omp parallel
#endif
	{
		int* w = new int[m];
		for (int i = 0; i < m; ++i)
			w[i] = -1;
#ifdef USE_OPENMP
#pragma omp for schedule(guided)
#endif
		for (int i = 0; i < n; ++i)
		{
			int r = Wp[i];
			int r0 = Wp[i];

			for (int k = Ap[i]; k < Ap[i + 1]; ++k)
			{
				int kcol = Ai[k];
				double kval = Av[k];

				for (int j = Bp[kcol]; j < Bp[kcol + 1]; ++j)
				{
					int jcol = Bi[j];
					double jval = Bv[j];

					if (w[jcol] < r0)
					{
						w[jcol] = r;
						Wi[r] = jcol;
						Wv[r++] = kval * jval;
					}
					else
						Wv[w[jcol]] += kval * jval;
				}
			}

			for (int k = Cp[i]; k < Cp[i + 1]; ++k)
			{
				int kcol = Ci[k];
				double kval = Cv[k];

				for (int j = Dp[kcol]; j < Dp[kcol + 1]; ++j)
				{
					int jcol = Di[j];
					double jval = Dv[j];
					
					if (w[jcol] < r0)
					{
						w[jcol] = r;
						Wi[r] = jcol;
						Wv[r++] = kval * jval;
					}
					else
						Wv[w[jcol]] += kval * jval;
				}
			}
		}

		delete[] w;
	}

	W.Free();
	W.size[0] = n;
	W.size[1] = m;
	W.rowptr = Wp;
	W.colind = Wi;
	W.values = Wv;
}

#define MPI_TAG 300

void ParCSR_MatMul(const ParCSR_Matrix& A, const ParCSR_Matrix& B, ParCSR_Matrix& C)
{
	MPI_Comm comm = A.comm;

	int rank;
	MPI_Comm_rank(comm, &rank);

	int num_neighbors = A.num_neighbors;
	int* neighbor_ranks = A.neighbor_ranks;

	int* recv_ptr = A.Ext_ptr;
	int* send_ptr = A.send_ptr;
	int* send_list = A.send_list;

	int n = A.LocLoc.size[0];
	int m = B.LocLoc.size[1];

	int* B_LocLoc_rowptr = B.LocLoc.rowptr;
	int* B_LocLoc_colind = B.LocLoc.colind;
	double* B_LocLoc_values = B.LocLoc.values;

	int* B_LocExt_rowptr = B.LocExt.rowptr;
	int* B_LocExt_colind = B.LocExt.colind;
	double* B_LocExt_values = B.LocExt.values;

	int B_num_neighbors = B.num_neighbors;
	int* B_neighbors = B.neighbor_ranks;
	int* B_Ext_ptr = B.Ext_ptr;
	int* B_Ext_map = B.Ext_map;

	int* B_Ext_from = new int[B.LocExt.size[1]];
	for (int r = 0; r < B_num_neighbors; ++r)
		for (int i = B_Ext_ptr[r]; i < B_Ext_ptr[r + 1]; ++i)
			B_Ext_from[i] = r;

	std::map<int, int> map_neighbor;
	for (int r = 0; r < B_num_neighbors; ++r)
		map_neighbor.insert(std::pair<int, int>(B_neighbors[r], r));
	for (int r = 0; r < num_neighbors; ++r)
		map_neighbor.insert(std::pair<int, int>(neighbor_ranks[r], map_neighbor.size()));

	int C_num_neighbors = map_neighbor.size();
	int* C_neighbors = new int[C_num_neighbors];
	for (std::map<int, int>::iterator iter = map_neighbor.begin(); iter != map_neighbor.end(); ++iter)
		C_neighbors[iter->second] = iter->first;

	std::set<int>* set_neighbors = new std::set<int>[C_num_neighbors];

	for (int r = 0; r < num_neighbors; ++r)
	{
		int rfrom = map_neighbor[neighbor_ranks[r]];
		for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
		{
			int irow = send_list[i];
			for (int j = B_LocExt_rowptr[irow]; j < B_LocExt_rowptr[irow + 1]; ++j)
			{
				int jfrom = B_Ext_from[B_LocExt_colind[j]];
				if (jfrom != rfrom)
				{
					set_neighbors[rfrom].insert(B_neighbors[jfrom]);
					set_neighbors[jfrom].insert(neighbor_ranks[r]);
				}
			}
		}
	}

	int** recv_c_nb = new int* [C_num_neighbors];
	int** send_c_nb = new int* [C_num_neighbors];

	MPI_Request* recv_requests = new MPI_Request[C_num_neighbors];
	MPI_Request* send_requests = new MPI_Request[C_num_neighbors];
	MPI_Status status;

	for (int r = 0; r < C_num_neighbors; ++r)
	{
		send_c_nb[r] = new int[set_neighbors[r].size()];
		int i = 0;
		for (std::set<int>::iterator iter = set_neighbors[r].begin(); iter != set_neighbors[r].end(); ++iter)
			send_c_nb[r][i++] = *iter;
		MPI_Isend(send_c_nb[r], set_neighbors[r].size(), MPI_INT, C_neighbors[r], MPI_TAG, comm, send_requests + r);
	}

	for (int r = 0; r < C_num_neighbors; ++r)
	{
		int cnt;
		MPI_Probe(C_neighbors[r], MPI_TAG, comm, &status);
		MPI_Get_count(&status, MPI_INT, &cnt);
		recv_c_nb[r] = new int[cnt];
		MPI_Recv(recv_c_nb[r], cnt, MPI_INT, C_neighbors[r], MPI_TAG, comm, &status);
		for (int i = 0; i < cnt; ++i)
			map_neighbor.insert(std::pair<int, int>(recv_c_nb[r][i], map_neighbor.size()));
		delete[] recv_c_nb[r];
	}

	for (int r = 0; r < C_num_neighbors; ++r)
	{
		MPI_Wait(send_requests + r, &status);
		delete[] send_c_nb[r];
	}

	delete[] set_neighbors;
	delete[] C_neighbors;
	delete[] recv_c_nb;
	delete[] send_c_nb;

	C_num_neighbors = map_neighbor.size();
	C_neighbors = new int[C_num_neighbors];
	for (std::map<int, int>::iterator iter = map_neighbor.begin(); iter != map_neighbor.end(); ++iter)
		C_neighbors[iter->second] = iter->first;

	int* send_bufsize = new int[num_neighbors];
	int* recv_bufsize = new int[num_neighbors];

	for (int r = 0; r < num_neighbors; ++r)
		if (recv_ptr[r + 1] > recv_ptr[r])
			MPI_Irecv(recv_bufsize + r, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			int num_rows = send_ptr[r + 1] - send_ptr[r];
			int num_nonzeros = 0;
			for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
				num_nonzeros += B_LocLoc_rowptr[send_list[i] + 1] - B_LocLoc_rowptr[send_list[i]] + B_LocExt_rowptr[send_list[i] + 1] - B_LocExt_rowptr[send_list[i]];
			send_bufsize[r] = (num_rows + 2 * num_nonzeros) * sizeof(int) + num_nonzeros * sizeof(double);
			MPI_Send(send_bufsize + r, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
		if (recv_ptr[r + 1] > recv_ptr[r])
			MPI_Wait(recv_requests + r, &status);

	char** recv_buffer = new char* [num_neighbors];
	char** send_buffer = new char* [num_neighbors];
	
	for (int r = 0; r < num_neighbors; ++r)
	{
		if (recv_ptr[r + 1] > recv_ptr[r])
		{
			recv_buffer[r] = (char*)malloc(recv_bufsize[r]);
			MPI_Irecv(recv_buffer[r], recv_bufsize[r], MPI_PACKED, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			send_buffer[r] = (char*)malloc(send_bufsize[r]);

			int position = 0;
			for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
			{
				int irow = send_list[i];
				int length = B_LocLoc_rowptr[irow + 1] - B_LocLoc_rowptr[irow] + B_LocExt_rowptr[irow + 1] - B_LocExt_rowptr[irow];

				int* colfrom = new int[length];
				int* colind = new int[length];
				double* values = new double[length];

				int j = 0;
				for (int k = B_LocLoc_rowptr[irow]; k < B_LocLoc_rowptr[irow + 1]; ++k)
				{
					colfrom[j] = rank;
					colind[j] = B_LocLoc_colind[k];
					values[j++] = B_LocLoc_values[k];
				}
				for (int k = B_LocExt_rowptr[irow]; k < B_LocExt_rowptr[irow + 1]; ++k)
				{
					int kcol = B_LocExt_colind[k];
					colfrom[j] = B_neighbors[B_Ext_from[kcol]];
					colind[j] = B_Ext_map[kcol];
					values[j++] = B_LocExt_values[k];
				}

				MPI_Pack(&length, 1, MPI_INT, send_buffer[r], send_bufsize[r], &position, comm);
				MPI_Pack(colfrom, length, MPI_INT, send_buffer[r], send_bufsize[r], &position, comm);
				MPI_Pack(colind, length, MPI_INT, send_buffer[r], send_bufsize[r], &position, comm);
				MPI_Pack(values, length, MPI_DOUBLE, send_buffer[r], send_bufsize[r], &position, comm);

				delete[] colfrom;
				delete[] colind;
				delete[] values;
			}

			MPI_Isend(send_buffer[r], send_bufsize[r], MPI_PACKED, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);
		}
	}

	std::map<int, int>* map_Ext = new std::map<int, int>[C_num_neighbors];
	for (int r = 0; r < B_num_neighbors; ++r)
		for (int i = B_Ext_ptr[r], j = 0; i < B_Ext_ptr[r + 1]; ++i, ++j)
			map_Ext[r].insert(std::pair<int, int>(B_Ext_map[i], j));

	int B_Ext_nnz = 0;
	for (int r = 0; r < num_neighbors; ++r)
	{
		if (recv_ptr[r + 1] > recv_ptr[r])
		{
			int num_rows = recv_ptr[r + 1] - recv_ptr[r];
			int num_nonzeros = (recv_bufsize[r] - num_rows * sizeof(int)) / (2 * sizeof(int) + sizeof(double));
			B_Ext_nnz += num_nonzeros;
		}
	}

	int* B_ExtLoc_rowptr = new int[A.LocExt.size[1] + 1];
	int* B_ExtLoc_colind = new int[B_Ext_nnz];
	double* B_ExtLoc_values = new double[B_Ext_nnz];
	int* B_ExtExt_rowptr = new int[A.LocExt.size[1] + 1];
	int* B_ExtExt_colind = new int[B_Ext_nnz];
	double* B_ExtExt_values = new double[B_Ext_nnz];
	int* B_ExtExt_colfrom = new int[B_Ext_nnz];

	B_ExtLoc_rowptr[0] = 0;
	B_ExtExt_rowptr[0] = 0;
	for (int r = 0; r < num_neighbors; ++r)
	{
		if (recv_ptr[r + 1] > recv_ptr[r])
		{
			MPI_Wait(recv_requests + r, &status);

			int position = 0;
			for (int i = recv_ptr[r], k = B_ExtLoc_rowptr[i], t = B_ExtExt_rowptr[i]; i < recv_ptr[r + 1]; B_ExtLoc_rowptr[++i] = k, B_ExtExt_rowptr[i] = t)
			{
				int length;
				MPI_Unpack(recv_buffer[r], recv_bufsize[r], &position, &length, 1, MPI_INT, comm);

				int* colfrom = new int[length];
				int* colind = new int[length];
				double* values = new double[length];

				MPI_Unpack(recv_buffer[r], recv_bufsize[r], &position, colfrom, length, MPI_INT, comm);
				MPI_Unpack(recv_buffer[r], recv_bufsize[r], &position, colind, length, MPI_INT, comm);
				MPI_Unpack(recv_buffer[r], recv_bufsize[r], &position, values, length, MPI_DOUBLE, comm);

				for (int j = 0; j < length; ++j)
				{
					if (colfrom[j] == rank)
					{
						B_ExtLoc_colind[k] = colind[j];
						B_ExtLoc_values[k++] = values[j];
					}
					else
					{
						int jfrom = map_neighbor[colfrom[j]];
						B_ExtExt_colfrom[t] = jfrom;
						std::pair<std::map<int, int>::iterator, bool> ret = map_Ext[jfrom].insert(std::pair<int, int>(colind[j], map_Ext[jfrom].size()));
						B_ExtExt_colind[t] = ret.first->second;
						B_ExtExt_values[t++] = values[j];
					}
				}

				delete[] colfrom;
				delete[] colind;
				delete[] values;
			}
			
			free(recv_buffer[r]);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			MPI_Wait(send_requests + r, &status);
			free(send_buffer[r]);
		}
	}

	delete[] recv_bufsize;
	delete[] recv_buffer;
	delete[] recv_requests;
	delete[] send_bufsize;
	delete[] send_buffer;
	delete[] send_requests;

	int* C_Ext_ptr = new int[C_num_neighbors + 1];
	C_Ext_ptr[0] = 0;
	for (int r = 0; r < C_num_neighbors; ++r)
		C_Ext_ptr[r + 1] = C_Ext_ptr[r] + map_Ext[r].size();

	int* C_Ext_map = new int[C_Ext_ptr[C_num_neighbors]];
	for (int r = 0; r < C_num_neighbors; ++r)
		for (std::map<int, int>::iterator iter = map_Ext[r].begin(); iter != map_Ext[r].end(); ++iter)
			C_Ext_map[C_Ext_ptr[r] + iter->second] = iter->first;

	for (int j = 0; j < B_ExtExt_rowptr[A.LocExt.size[1]]; ++j)
		B_ExtExt_colind[j] += C_Ext_ptr[B_ExtExt_colfrom[j]];

	int* _B_LocExt_colind = new int[B_LocExt_rowptr[B.LocExt.size[0]]];

	for (int j = 0; j < B_LocExt_rowptr[B.LocExt.size[0]]; ++j)
	{
		int r = B_Ext_from[B_LocExt_colind[j]];
		_B_LocExt_colind[j] = B_LocExt_colind[j] - B_Ext_ptr[r] + C_Ext_ptr[r];
	}

	delete[] map_Ext;
	delete[] B_Ext_from;
	delete[] B_ExtExt_colfrom;

	CSR_Matrix A_LocLoc, A_LocExt;

	A_LocLoc.Reference(A.LocLoc);
	A_LocExt.Reference(A.LocExt);

	CSR_Matrix B_LocLoc, B_LocExt;

	B_LocLoc.Reference(B.LocLoc);
	B_LocExt.size[0] = B_LocExt.size[0];
	B_LocExt.size[1] = C_Ext_ptr[C_num_neighbors];
	B_LocExt.rowptr = B.LocExt.rowptr;
	B_LocExt.colind = _B_LocExt_colind;
	B_LocExt.values = B.LocExt.values;
	B_LocExt.ref = 1;

	CSR_Matrix B_ExtLoc, B_ExtExt;

	B_ExtLoc.size[0] = A.LocExt.size[1];
	B_ExtLoc.size[1] = m;
	B_ExtLoc.rowptr = B_ExtLoc_rowptr;
	B_ExtLoc.colind = B_ExtLoc_colind;
	B_ExtLoc.values = B_ExtLoc_values;

	B_ExtExt.size[0] = A.LocExt.size[1];
	B_ExtExt.size[1] = C_Ext_ptr[C_num_neighbors];
	B_ExtExt.rowptr = B_ExtExt_rowptr;
	B_ExtExt.colind = B_ExtExt_colind;
	B_ExtExt.values = B_ExtExt_values;

	CSR_Matrix C_LocLoc, C_LocExt;

	CSR_MulAddMul(A_LocLoc, B_LocLoc, A_LocExt, B_ExtLoc, C_LocLoc);
	CSR_MulAddMul(A_LocLoc, B_LocExt, A_LocExt, B_ExtExt, C_LocExt);

	C_LocLoc.ref = 1;
	C_LocExt.ref = 1;

	delete[] _B_LocExt_colind;

	int* C_LocExt_rowptr = C_LocExt.rowptr;
	int* C_LocExt_colind = C_LocExt.colind;

	int* v = new int[C_Ext_ptr[C_num_neighbors]];
	for (int i = 0; i < C_Ext_ptr[C_num_neighbors]; ++i)
		v[i] = -1;

	for (int j = C_LocExt_rowptr[0]; j < C_LocExt_rowptr[n]; ++j)
		v[C_LocExt_colind[j]] = 1;

	for (int r = 0, k = 0, i = 0; r < C_num_neighbors; C_Ext_ptr[++r] = k)
	{
		for ( ; i < C_Ext_ptr[r + 1]; ++i)
		{
			if (v[i] == 1)
			{
				v[i] = k;
				C_Ext_map[k++] = C_Ext_map[i];
			}
		}
	}

	for (int j = C_LocExt_rowptr[0]; j < C_LocExt_rowptr[n]; ++j)
		C_LocExt_colind[j] = v[C_LocExt_colind[j]];

	C_LocExt.size[1] = C_Ext_ptr[C_num_neighbors];

	delete[] v;

	C.Free();

	C.comm = comm;

	C.LocLoc.size[0] = C_LocLoc.size[0];
	C.LocLoc.size[1] = C_LocLoc.size[1];
	C.LocLoc.rowptr = C_LocLoc.rowptr;
	C.LocLoc.colind = C_LocLoc.colind;
	C.LocLoc.values = C_LocLoc.values;

	C.LocExt.size[0] = C_LocExt.size[0];
	C.LocExt.size[1] = C_LocExt.size[1];
	C.LocExt.rowptr = C_LocExt.rowptr;
	C.LocExt.colind = C_LocExt.colind;
	C.LocExt.values = C_LocExt.values;

	C.num_neighbors = C_num_neighbors;
	C.neighbor_ranks = C_neighbors;
	C.Ext_ptr = C_Ext_ptr;
	C.Ext_map = C_Ext_map;
}

}